import bentoml
import numpy as np

@bentoml.service
class IrisClassifier:
    def __init__(self):
        self.iris_clf = bentoml.sklearn.load_model("iris_clf:latest")

    @bentoml.api
    def classify(self, input_data: np.ndarray) -> dict:
        """
        Classify iris flower based on input features.
        
        Args:
            input_data: A 2D array of shape (n_samples, 4) where each row represents
                       [sepal_length, sepal_width, petal_length, petal_width]

                [   [5.1, 3.5, 1.4, 0.2],
                    [4.9, 3.0, 1.4, 0.2],
                    [6.4, 3.2, 4.5, 1.5],
                    [6.3, 3.3, 6.0, 2.5],
                    [7.7, 3.8, 6.7, 2.2]
                ]

        Returns:
            dict: Prediction results with class labels
        """
        prediction = self.iris_clf.predict(input_data)
        return {"prediction": prediction.tolist()} 